#include <iostream>
#include <omp.h>
#include <vector>
#include <cstdlib>
#include <string>
#include <fstream>
#include <sys/time.h>
#include <mpi.h>
#include <pmmintrin.h>
#include <cmath>
#define millitime(x) (x.tv_sec * 1000 + x.tv_usec / 1000.0)
#define MaxN 2048
using namespace std;

int n, thread_count;
string path;
string file_num = "1";
float **A = NULL, **mother = NULL;

void common_gauss() {
    for (int k = 0; k < n; k++) {
        float ele = A[k][k];
        for (int j = k + 1; j < n; j++)
            A[k][j] = A[k][j] / ele;
        A[k][k] = 1.0;
        for (int i = k + 1; i < n; i++) {
            for (int j = k + 1; j < n; j++)
                A[i][j] = A[i][j] - A[i][k] * A[k][j];
            A[i][k] = 0;
        }
    }
}

void init() {
    A = new float*[n];
    mother = new float*[n];
    for (int i = 0; i < n; i++) {
        A[i] = new float[n];
        mother[i] = new float[n];
        for (int j = 0; j < n; j++) {
            A[i][j] = 0;
            mother[i][j] = 0;
        }
    }
    for (int i = 0; i < n; i++)
        for (int j = i; j < n; j++)
            mother[i][j] = (j == i) ? 1 : i + j;
}

void release_matrix() {
    for (int i = 0; i < n; i++) {
        delete[] A[i];
        delete[] mother[i];
    }
    delete[] A;
    delete[] mother;
}

void arr_reset() {
    for (int i = 0; i < n; i++)
        for (int j = 0; j < n; j++)
            A[i][j] = mother[i][j];
    for (int i = 1; i < n; i++)
        for (int j = 0; j < i; j++)
            for (int k = 0; k < n; k++)
                A[i][k] += mother[j][k];
}

void printResult() {
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            cout << mother[i][j] << ' ';
        }
        cout << endl;
    }
    cout << endl;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            cout << A[i][j] << ' ';
        }
        cout << endl;
    }
}

void testResult() {
    for (int i = 0; i < n; i++)
        for (int j = 0; j < n; j++)
            if (abs(A[i][j] - mother[i][j]) >= 1e-5) {
                cout << "Something wrong!" << endl;
                cout << i << ' ' << j << ' ' << A[i][j] << ' ' << mother[i][j] << endl;
                exit(-1);
            }
}

void block_gauss_opt(int my_rank, int num_proc) {   
    __m128 v0, v1, v2;
    int block_size = n / num_proc;
    int remain = n % num_proc;

    int my_begin = my_rank * block_size;
    int my_end = my_rank == num_proc - 1 ? my_begin + block_size + remain : my_begin + block_size;
    int k, j, i;
    #pragma omp parallel num_threads(thread_count), private(v0, v1, v2, k, j, i)
    for (k = 0; k < n; k++) {
        #pragma omp single 
        {
            if (k >= my_begin && k < my_end) {
                v1 = _mm_set_ps(A[k][k], A[k][k], A[k][k], A[k][k]);
                for (j = k + 1; j <= n - 4; j += 4) {
                    v0 = _mm_loadu_ps(A[k] + j);
                    v0 = _mm_div_ps(v0, v1);
                    _mm_storeu_ps(A[k] + j, v0);
                }
                float ele = A[k][k];
                for (j; j < n; j++)
                    A[k][j] = A[k][j] / ele;
                A[k][k] = 1.0;
                for (j = my_rank + 1; j < num_proc; j++)
                    MPI_Send(A[k], n, MPI_FLOAT, j, 2, MPI_COMM_WORLD);
            }
            else {
                int current_work_p = k / block_size;
                if (current_work_p < my_rank)
                    MPI_Recv(A[k], n, MPI_FLOAT, current_work_p, 2, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            }
        }
        #pragma omp for
        for (i = my_begin; i < my_end; i++) {
            if (i <= k) 
                continue;
            v1 = _mm_set_ps(A[i][k], A[i][k], A[i][k], A[i][k]);
            for (j = k + 1; j <= n - 4; j += 4) {
                v2 = _mm_loadu_ps(A[k] + j);
                v0 = _mm_loadu_ps(A[i] + j);
                v2 = _mm_mul_ps(v1, v2);
                v0 = _mm_sub_ps(v0, v2);
                _mm_storeu_ps(A[i] + j, v0);
            }
            for (j; j < n; j++)
                A[i][j] = A[i][j] - A[i][k] * A[k][j]; 
            A[i][k] = 0.0;
        }
    }
}

int main(int argc, char *argv[])
{
    thread_count = atoi(argv[2]);
    n = atoi(argv[1]);
    MPI_Init(NULL, NULL);
    init();
    MPI_Barrier(MPI_COMM_WORLD);
    block_gauss_opt(0, 1);
    MPI_Barrier(MPI_COMM_WORLD);
    release_matrix();
    MPI_Finalize();
    return 0;
}
